import numpy as np
from scipy.spatial.transform import Rotation as R


def matrix_to_euler(matrix, order='xyz'):
    """
    Convert a 4x4 transformation matrix to Euler angles.

    Parameters:
    matrix (numpy.ndarray): A 4x4 transformation matrix.
    order (str): The order of Euler angles (e.g., 'rpy', 'xyz', etc.).

    Returns:
    tuple: Euler angles (roll, pitch, yaw) in radians.
    """
    # Extract the rotation matrix from the transformation matrix
    rotation_matrix = matrix[:3, :3]

    # Convert the rotation matrix to an Euler angle
    euler = R.from_matrix(rotation_matrix).as_euler(order, degrees=False)

    return euler


def matrix_to_translation(matrix):
    """
    Extract the translation vector from a 4x4 transformation matrix.

    Parameters:
    matrix (numpy.ndarray): A 4x4 transformation matrix.

    Returns:
    numpy.ndarray: The translation vector (x, y, z).
    """
    return matrix[:3, 3]


# Example 4x4 transformation matrix
matrix = np.array([
        [
          0.0107962,
          -0.999492,
          -0.0299938,
          0
        ],
        [
          0.999942,
          0.0107913,
          0.000323837,
          1.3
        ],
        [
          0,
          -0.0299955,
          0.99955,
          1.35
        ],
        [
          0.0,
          0.0,
          0.0,
          1.0
        ]
      ])

def mat2xyzrpy(matrix,order='xyz'):
    # For carla
    euler_angles = matrix_to_euler(matrix, order=order)
    translation = matrix_to_translation(matrix)
    spawn_point = {
        "x": translation[0],
        "y": translation[1],
        "z": translation[2],
        "roll": euler_angles[0]*180.0/3.14159,
        "pitch": euler_angles[1]*180.0/3.14159,
        "yaw": euler_angles[2]*180.0/3.14159
    }
    # print(spawn_point)
    print(json.dumps(spawn_point, indent=4))


if __name__ == '__main__':
    import json
    calib_path = '/media/adt/ZWH4T/ZWH/Dataset3d/final/cyw2/__2023-05-28-17-23-20_filter11-cars/calib/calib.json'
    with open(calib_path) as f:
        calib_data = json.load(f)

    sensor_names = ['lidar1', 'lidar2', 'lidar3', 'lidar4', 'radar3']
    camera_names = ['camera75', 'camera77', 'camera80', 'camera81']
    for lidar_name in sensor_names:
        dar2base = np.array(calib_data[lidar_name]['transforms']['base_link'])
        print(lidar_name)
        mat2xyzrpy(dar2base,order='xyz')

    for camera_name in camera_names:
        cam2base = np.array(calib_data[camera_name]['transforms']['base_link'])
        cam2xyz = np.array([[0,0,1,0],
                            [-1,0,0,0],
                            [0,-1,0,0],
                            [0,0,0,1]])
        xyz2base = cam2base @ np.linalg.inv(cam2xyz)
        print(camera_name)
        mat2xyzrpy(xyz2base,order='xyz')

    pass
